{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Evolutive equilibrium calculations\n",
    "\n",
    "This notebook will demonstrate the core use-case for FreeGSNKE: simulating the (time-dependent) evolution of Grad-Shafranov (GS) equilibria. In particular we will simulate a **vertical displacement event (VDE)**.\n",
    "\n",
    "To do this, we need to:\n",
    "- build the tokamak machine.\n",
    "- instatiate a GS equilibrium (to be used as an initial condition for the evolutive solver).\n",
    "- calculate a vertical instability growth rate for this equilibrium and carry out passive structure mode removal via a normal mode decomposition (i.e. removing modes that have little effect on the evolution). \n",
    "- define time-dependent plasma current density profile parameters and coil voltages.\n",
    "- evolve the active coil currents, the total plasma current, and the equilbirium using these profile parameters and voltages by solving the circuit equations alongside the GS equation.\n",
    "\n",
    "Refer to the paper by [Amorisco et al. (2024)](https://pubs.aip.org/aip/pop/article/31/4/042517/3286904/FreeGSNKE-A-Python-based-dynamic-free-boundary) for more details. \n",
    "\n",
    "We should note that here we will use **fixed** (time-independent) profile parameters and voltages to simulate a VDE, however, truly time-dependent parameters would be required to simulate a plasma shot (see future notebooks). \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from copy import deepcopy\n",
    "from IPython.display import display, clear_output\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the machine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build machine\n",
    "from freegsnke import build_machine\n",
    "tokamak = build_machine.tokamak(\n",
    "    active_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_active_coils.pickle\",\n",
    "    passive_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_passive_coils.pickle\",\n",
    "    limiter_path=f\"../machine_configs/MAST-U/MAST-U_like_limiter.pickle\",\n",
    "    wall_path=f\"../machine_configs/MAST-U/MAST-U_like_wall.pickle\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Diverted plasma example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate an equilibrium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import equilibrium_update\n",
    "\n",
    "eq = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak,\n",
    "    Rmin=0.1, Rmax=2.0,   # Radial range\n",
    "    Zmin=-2.2, Zmax=2.2,  # Vertical range\n",
    "    nx=65,                # Number of grid points in the radial direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    ny=129,               # Number of grid points in the vertical direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    # psi=plasma_psi\n",
    ")  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instatiate a profile object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke.jtor_update import ConstrainPaxisIp\n",
    "\n",
    "profiles = ConstrainPaxisIp(\n",
    "    eq=eq,        # equilibrium object\n",
    "    paxis=8.1e3,  # profile object\n",
    "    Ip=6.2e5,     # plasma current\n",
    "    fvac=0.5,     # fvac = rB_{tor}\n",
    "    alpha_m=1.8,  # profile function parameter\n",
    "    alpha_n=1.2   # profile function parameter\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set coil currents\n",
    "Here we set coil currents that create a diverted plasma (as seen in prior notebooks). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('simple_diverted_currents_PaxisIp.pk', 'rb') as f:\n",
    "    current_values = pickle.load(f)\n",
    "\n",
    "for key in current_values.keys():\n",
    "    eq.tokamak.set_coil_current(coil_label=key, current_value=current_values[key])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instatiate the solver\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import GSstaticsolver\n",
    "GSStaticSolver = GSstaticsolver.NKGSsolver(eq)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Call forward solver to find equilibrium "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GSStaticSolver.solve(eq=eq, \n",
    "                     profiles=profiles, \n",
    "                     constrain=None, \n",
    "                     target_relative_tolerance=1e-8,\n",
    "                     verbose=0\n",
    "                     )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the initial equilibrium "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq.plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Time evolution\n",
    "\n",
    "We are now ready to solve the forward time-evolutive problem.\n",
    "\n",
    "#### Initialise nonlinear (evolutive) solver\n",
    "To start, we need to initialise the evolutive solver object, `freegsnke.nonlinear_solve.nl_solver`.\n",
    "\n",
    "The time-evolutive solver requires:\n",
    "- `eq`: an equilibrium to inform the solver of the machine and domain properties.\n",
    "- `profiles`: defined above.\n",
    "- `full_timestep`: the time step by which time is advanced at every call of the stepper (this can be modified later depending on the growth rate of the equilibrium). An appropriate time step can also be set based on the growth rate calculation. Use `automatic_timestep` to set the time step in units of the (inverse of the) growth rate.\n",
    "- `plasma_resistivity`: resistivity of the plasma (which here is assumed to be constant during the time evolution but can be made time-dependent if known).\n",
    "- `min_dIy_dI`: threshold value below which passive structure normal modes are dropped. Modes with norm(d(Iy)/dI)<`min_dIy_dI` are dropped, which filters out modes that do not actually couple with the plasma.\n",
    "- `max_mode_frequency`: threshold value for characteristic frequencies above which passive structure normal modes are dropped (i.e. the fast modes).\n",
    "\n",
    "Other customisable inputs are available, do see the documentation for more details. For example, one may explicitly set your own resistance and inductance matries for the tokamak machine. \n",
    "\n",
    "The solver can be used on different equilibria and/or profiles, but these need to have the same machine, domain, and limiter as the one used at the time of the solver instantiation. For different machines, a new time-evolutive solver should be created.\n",
    "\n",
    "The input equilibrium and profile functions are also used as the expansion point around which the dynamics are linearised."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import nonlinear_solve\n",
    "\n",
    "stepping = nonlinear_solve.nl_solver(\n",
    "    eq=eq, \n",
    "    profiles=profiles, \n",
    "    full_timestep=5e-4, \n",
    "    plasma_resistivity=1e-6,\n",
    "    min_dIy_dI=0.1,\n",
    "    max_mode_frequency=10**2.5,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set active coil voltages\n",
    "\n",
    "In this example, we will evolve a plasma in absence of any control policy or current drive.\n",
    "\n",
    "Just as an example, the following calculates active voltages to be applied to the poloidal field coils (and Solenoid) using $V = RI$, with current values as defined by the initial equilibrium (i.e. we have constant voltages).\n",
    "\n",
    "In most FreeGSNKE use cases, these active voltages will be determined by a control policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "voltages = (stepping.vessel_currents_vec*stepping.evol_metal_curr.R)[:stepping.evol_metal_curr.n_active_coils] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start, the solver is prepared by setting the initial conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stepping.initialize_from_ICs(eq, profiles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set time steps and storage lists\n",
    "Now we set the total number of time steps we want to simulate and initialise some variables for logging the evolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of time steps to simulate\n",
    "max_count = 45\n",
    "\n",
    "# initialising some variables for iteration and logging\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "history_times = [t]\n",
    "history_currents = [stepping.currents_vec]\n",
    "history_equilibria = [deepcopy(stepping.eq1)]\n",
    "separatrix = stepping.eq1.separatrix(ntheta=100)\n",
    "history_width = [np.amax(separatrix[:,0]) - np.amin(separatrix[:,0])]\n",
    "history_o_points = [stepping.eq1.opt[0]]\n",
    "history_elongation = [(np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width[0]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call the solver (linear)\n",
    "Finally, we call the time-evolutive solver with `stepping.nlstepper()` sequentially until we reach the preset end time.\n",
    "\n",
    "The following demonstrates a solely linear evolution of the plasma by setting `linear_only=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loop over time steps\n",
    "while counter<max_count:\n",
    "    clear_output(wait=True)\n",
    "    display(f'Step: {counter}/{max_count-1}')\n",
    "    display(f'Time: t = {t:.2e}')\n",
    "\n",
    "    # carry out the time step\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec=voltages, \n",
    "        linear_only=True,\n",
    "        verbose=0\n",
    "    )\n",
    "\n",
    "    # store information on the time step\n",
    "    t += stepping.dt_step\n",
    "    history_times.append(t)\n",
    "    counter += 1\n",
    "\n",
    "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
    "    history_currents.append(stepping.currents_vec)\n",
    "    history_equilibria.append(deepcopy(stepping.eq1))\n",
    "    separatrix = stepping.eq1.separatrix(ntheta=100)\n",
    "    history_width.append(np.amax(separatrix[:,0]) - np.amin(separatrix[:,0]))\n",
    "    history_o_points = np.append(history_o_points, [stepping.eq1.opt[0]], axis=0)\n",
    "    history_elongation.append((np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width[-1])\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents = np.array(history_currents)\n",
    "history_times = np.array(history_times)\n",
    "history_o_points = np.array(history_o_points)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "\n",
    "ax1.grid(True, which='both')\n",
    "history_equilibria[-1].plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call the solver (nonlinear)\n",
    "\n",
    "Now we evolve the same plasma according to the full nonlinear dynamics for the same time interval. This is done by using `linear_only=False` in the call to the stepper.\n",
    "\n",
    "We need to re-initialise from the initial conditions and reset the counter, but otherwise the method is identical to the one above.\n",
    "\n",
    "Note that the full nonlinear evolutive solve is a lot more computationally expensive than solving the linear evolutive problem. As such, the following cell may take several minutes to execute, depending on your hardware."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# reset the solver object by resetting the initial conditions\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "# initialising some variables for iteration and logging\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "history_times_nl = [t]\n",
    "history_currents_nl = [stepping.currents_vec]\n",
    "history_equilibria_nl = [deepcopy(stepping.eq1)]\n",
    "separatrix = stepping.eq1.separatrix(ntheta=100)\n",
    "history_width_nl = [np.amax(separatrix[:,0]) - np.amin(separatrix[:,0])]\n",
    "history_o_points_nl = [stepping.eq1.opt[0]]\n",
    "history_elongation_nl = [(np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width[0]]\n",
    "\n",
    "# loop over the time steps\n",
    "while counter<max_count:\n",
    "    clear_output(wait=True)\n",
    "    display(f'Step: {counter}/{max_count-1}')\n",
    "    display(f'Time: t = {t:.2e}')\n",
    "    \n",
    "    # carry out the time step\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec=voltages, \n",
    "        linear_only=False,\n",
    "        verbose=0\n",
    "    )\n",
    "\n",
    "    # store information on the time step\n",
    "    t += stepping.dt_step\n",
    "    history_times_nl.append(t)\n",
    "    counter += 1\n",
    "    \n",
    "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
    "    history_currents_nl.append(stepping.currents_vec)\n",
    "    history_equilibria_nl.append(deepcopy(stepping.eq1))\n",
    "    separatrix = stepping.eq1.separatrix(ntheta=100)\n",
    "    history_width_nl.append(np.amax(separatrix[:,0]) - np.amin(separatrix[:,0]))\n",
    "    history_o_points_nl = np.append(history_o_points_nl, [stepping.eq1.opt[0]], axis=0)\n",
    "    history_elongation_nl.append((np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width_nl[-1])\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents_nl = np.array(history_currents_nl)\n",
    "history_times_nl = np.array(history_times_nl)\n",
    "history_o_points_nl = np.array(history_o_points_nl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot some time-evolving quantities\n",
    "The following plots the evolution of a number of tracked values and compares the linear/nonlinear evoltuions against one another. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot evolution of tracked values and compare between linear and non-linear evolution\n",
    "fig, axs = plt.subplots(2, 3, figsize=(10, 5), dpi=80, constrained_layout=True)\n",
    "axs_flat = axs.flat\n",
    "\n",
    "axs_flat[0].plot(history_times, history_o_points[:, 0],'k+', label='linear')\n",
    "axs_flat[0].plot(history_times_nl, history_o_points_nl[:, 0],'rx', label='nonlinear')\n",
    "axs_flat[0].set_xlabel('Time')\n",
    "axs_flat[0].set_ylabel('O-point $R$')\n",
    "axs_flat[0].legend()\n",
    "\n",
    "axs_flat[1].plot(history_times, abs(history_o_points[:, 1]),'k+')\n",
    "axs_flat[1].plot(history_times_nl, abs(history_o_points_nl[:, 1]),'rx')\n",
    "axs_flat[1].set_yscale('log')\n",
    "axs_flat[1].set_xlabel('Time')\n",
    "axs_flat[1].set_ylabel('abs( O-point $Z$ )')\n",
    "\n",
    "axs_flat[2].plot(history_times, history_o_points[:, 2],'k+')\n",
    "axs_flat[2].plot(history_times_nl, history_o_points_nl[:, 2],'rx')\n",
    "axs_flat[2].set_xlabel('Time')\n",
    "axs_flat[2].set_ylabel('O-point $\\Psi$')\n",
    "\n",
    "axs_flat[3].plot(history_times, history_currents[:,-1]*stepping.plasma_norm_factor,'k+')\n",
    "axs_flat[3].plot(history_times_nl, history_currents_nl[:,-1]*stepping.plasma_norm_factor,'rx')\n",
    "axs_flat[3].set_xlabel('Time')\n",
    "axs_flat[3].set_ylabel('Plasma current')\n",
    "\n",
    "axs_flat[4].plot(history_times, history_width,'k+')\n",
    "axs_flat[4].plot(history_times_nl, history_width_nl,'rx')\n",
    "axs_flat[4].set_xlabel('Time')\n",
    "axs_flat[4].set_ylabel('Plasma width')\n",
    "\n",
    "axs_flat[5].plot(history_times, history_elongation,'k+')\n",
    "axs_flat[5].plot(history_times_nl, history_elongation_nl,'rx')\n",
    "axs_flat[5].set_xlabel('Time')\n",
    "axs_flat[5].set_ylabel('Plasma elongation')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the equilibria at the final time step\n",
    "fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8), dpi=80)\n",
    "\n",
    "ax1.grid(True, which='both')\n",
    "history_equilibria[-1].plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "ax1.set_title(\"Linear solve\")\n",
    "\n",
    "ax2.grid(True, which='both')\n",
    "history_equilibria_nl[-1].plot(axis=ax2, show=False)\n",
    "eq.tokamak.plot(axis=ax2, show=False)\n",
    "ax2.set_xlim(0.1, 2.15)\n",
    "ax2.set_ylim(-2.25, 2.25)\n",
    "ax2.set_title(\"Nonlinear solve\")\n",
    "\n",
    "plt.tight_layout()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Limited plasma example\n",
    "\n",
    "Finally, we examine a limiter configuration (there is no fundamental difference from the process described above).\n",
    "\n",
    "First we re-instantiate the equilibrium object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eq = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak,\n",
    "    Rmin=0.1, Rmax=2.0,\n",
    "    Zmin=-2.2, Zmax=2.2,\n",
    "    nx=65, \n",
    "    ny=129\n",
    ")  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We load a set of coil currents that will give us a limiter plasma in our current tokamak."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "with open('limiter_currents.json', 'r') as f:\n",
    "    current_values = json.load(f)\n",
    "\n",
    "for key in current_values.keys():\n",
    "    eq.tokamak.set_coil_current(key, current_values[key])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we set the profiles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles = ConstrainPaxisIp(\n",
    "    eq=eq,\n",
    "    paxis=8.1e3,\n",
    "    Ip=6.2e5,\n",
    "    fvac=0.5,\n",
    "    alpha_m=1.8,\n",
    "    alpha_n=1.2\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final part of the setup is to do a forward solve and visualise the limiter equilibrium.\n",
    "The red dashed line is the flux surface through the first X-point, however, the actual last closed flux surface in limiter equilibria is displayed as a full black line. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GSStaticSolver.solve(eq, profiles, target_relative_tolerance=1e-8)\n",
    "\n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq.plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can do the nonlinear evolutive solve as before. First, we re-initialize the solver, which calculates the linearization of the dynamics relevant to this new equilibrium."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import nonlinear_solve\n",
    "\n",
    "stepping = nonlinear_solve.nl_solver(\n",
    "    eq=eq, profiles=profiles, \n",
    "    full_timestep=.5e-3, \n",
    "    plasma_resistivity=1e-6,    \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# recalculate the active voltages using the new currents\n",
    "U_active = (stepping.vessel_currents_vec*stepping.evol_metal_curr.R)[:stepping.evol_metal_curr.n_active_coils] \n",
    "\n",
    "# number of time steps to simulate\n",
    "max_count = 45\n",
    "\n",
    "# reset the solver object by resetting the initial conditions\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "# initialising some variables for iteration and logging\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "history_times_nl = [t]\n",
    "history_currents_nl = [stepping.currents_vec]\n",
    "history_equilibria_nl = [deepcopy(stepping.eq1)]\n",
    "separatrix = stepping.eq1.separatrix(ntheta=100)\n",
    "history_width_nl = [np.amax(separatrix[:,0]) - np.amin(separatrix[:,0])]\n",
    "history_o_points_nl = [stepping.eq1.opt[0]]\n",
    "history_elongation_nl = [(np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width[0]]\n",
    "\n",
    "# Simulate the forward evolution\n",
    "while counter<max_count:\n",
    "    clear_output(wait=True)\n",
    "    display(f'Step: {counter}/{max_count-1}')\n",
    "    display(f'current time t = {t:.2e}')\n",
    "    \n",
    "    # carry out the time step\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec=U_active, \n",
    "        linear_only=False,\n",
    "        target_relative_tol_currents=0.01,\n",
    "        target_relative_tol_GS=0.01,\n",
    "\n",
    "    )\n",
    "\n",
    "    # store information on the time step\n",
    "    t += stepping.dt_step\n",
    "    history_times_nl.append(t)\n",
    "    counter += 1\n",
    "    \n",
    "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
    "    history_currents_nl.append(stepping.currents_vec)\n",
    "    history_equilibria_nl.append(deepcopy(stepping.eq1))\n",
    "    separatrix = stepping.eq1.separatrix(ntheta=100)\n",
    "    history_width_nl.append(np.amax(separatrix[:,0]) - np.amin(separatrix[:,0]))\n",
    "    history_o_points_nl = np.append(history_o_points_nl, [stepping.eq1.opt[0]], axis=0)\n",
    "    history_elongation_nl.append((np.amax(separatrix[:,1]) - np.amin(separatrix[:,1]))/history_width_nl[-1])\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents_nl = np.array(history_currents_nl)\n",
    "history_times_nl = np.array(history_times_nl)\n",
    "history_o_points_nl = np.array(history_o_points_nl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot evolution of tracked values\n",
    "fig, axs = plt.subplots(2, 3, figsize=(10, 5), dpi=80, constrained_layout=True)\n",
    "axs_flat = axs.flat\n",
    "\n",
    "axs_flat[0].plot(history_times_nl, history_o_points_nl[:, 0],'rx', label='nonlinear')\n",
    "axs_flat[0].set_xlabel('Time')\n",
    "axs_flat[0].set_ylabel('O-point $R$')\n",
    "axs_flat[0].legend()\n",
    "\n",
    "axs_flat[1].plot(history_times_nl, abs(history_o_points_nl[:, 1]),'rx')\n",
    "axs_flat[1].set_yscale('log')\n",
    "axs_flat[1].set_xlabel('Time')\n",
    "axs_flat[1].set_ylabel('abs( O-point $Z$ )')\n",
    "\n",
    "axs_flat[2].plot(history_times_nl, history_o_points_nl[:, 2],'rx')\n",
    "axs_flat[2].set_xlabel('Time')\n",
    "axs_flat[2].set_ylabel('O-point $\\Psi$')\n",
    "\n",
    "axs_flat[3].plot(history_times_nl, history_currents_nl[:,-1]*stepping.plasma_norm_factor,'rx')\n",
    "axs_flat[3].set_xlabel('Time')\n",
    "axs_flat[3].set_ylabel('Plasma current')\n",
    "\n",
    "axs_flat[4].plot(history_times_nl, history_width_nl,'rx')\n",
    "axs_flat[4].set_xlabel('Time')\n",
    "axs_flat[4].set_ylabel('Plasma width')\n",
    "\n",
    "axs_flat[5].plot(history_times_nl, history_elongation_nl,'rx')\n",
    "axs_flat[5].set_xlabel('Time')\n",
    "axs_flat[5].set_ylabel('Plasma elongation')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the equilibria at the final time step\n",
    "fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8), dpi=80)\n",
    "\n",
    "ax1.grid(True, which='both')\n",
    "history_equilibria_nl[-1].plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "ax1.set_title('t=0')\n",
    "\n",
    "ax2.grid(True, which='both')\n",
    "history_equilibria_nl[-1].plot(axis=ax2, show=False)\n",
    "eq.tokamak.plot(axis=ax2, show=False)\n",
    "ax2.set_xlim(0.1, 2.15)\n",
    "ax2.set_ylim(-2.25, 2.25)\n",
    "ax2.set_title('t='+str(t))\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "freegsnke4e",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}